#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""计算成本表面
给定坐标点和成本栅格, 计算成本栅格和回溯栅格。
其中 -t 参数指定了成本栅格保存的数据类型, 包括: byte, int16, int32, float32, float64。默认 float32。

Example:
    $ python cost_surface.py -t float64 point.shp cost.tif
"""

import sys
from osgeo import gdal, ogr, osr
from skimage.graph import MCP_Geometric
import numpy as np

def vector2array(shape_fn, raster_fn):
    """将 Shapefile 文件按栅格坐标读入

    参数:
        shape_fn (str) Shapefile 文件路径

    返回:
        array 数组
    """
    raster = gdal.Open(raster_fn)
    band = raster.GetRasterBand(1)

    shape = ogr.Open(shape_fn)
    layer = shape.GetLayer()

    # 将矢量数据空间坐标系投影到栅格
    source_sr = layer.GetSpatialRef()
    target_sr = osr.SpatialReference()
    target_sr.ImportFromWkt(raster.GetProjectionRef())
    coord_trans = osr.CoordinateTransformation(source_sr, target_sr)

    points_x = []
    points_y = []

    for f_id in range(layer.GetFeatureCount()):
        feature = layer.GetFeature(f_id)

        geom = feature.GetGeometryRef()
        geom.Transform(coord_trans)

        if geom.GetGeometryName() == 'POINT':
            points = geom.GetPointCount()

            for p in range(points):
                lon, lat, z = geom.GetPoint(p)
                points_x.append(lon)
                points_y.append(lat)

        else:
            raise ValueError('unsupport geometry type')
            sys.exit(1)

    return zip(points_x, points_y)

def raster2array(raster_fn):
    """将 GeoTiff 文件以数组的形式读入

    参数:
        raster_fn (str) GeoTiff 文件路径

    返回:
        array 数组
    """
    raster = gdal.Open(raster_fn)
    band = raster.GetRasterBand(1)
    array = band.ReadAsArray()

    return array

def coord2pixel_offset(raster_fn, x, y):
    """已知空间坐标点 x, y 求其在 GeoTiff 文件数组中对应的索引

    参数:
        raster_fn (str) GeoTiff 文件路径
        x (float) x 坐标
        y (float) y 坐标

    返回:
        x_offset (int) x 偏移
        y_offset (int) y 偏移
    """
    raster = gdal.Open(raster_fn)
    geotransform = raster.GetGeoTransform()
    origin_x = geotransform[0]
    origin_y = geotransform[3]
    pixel_width = geotransform[1]
    pixel_heigh = geotransform[5]
    x_offset = int((x - origin_x) / pixel_width)
    y_offset = int((y - origin_y) / pixel_heigh)

    # print(F"x: {x}, y: {y}, x_offset: {x_offset}, y_offset: {y_offset}")

    return x_offset, y_offset

def setNoDataValue(raster_fn, value_array, data_type):
    raster = gdal.Open(raster_fn)
    band = raster.GetRasterBand(1)
    raster_array = band.ReadAsArray()
    raster_nodata = band.GetNoDataValue()

    if data_type == gdal.GDT_Byte:
        no_data = 255
    elif data_type == gdal.GDT_Int16:
        no_data = -32768
    elif data_type == gdal.GDT_Int32:
        no_data = -2147483648
    elif data_type == gdal.GDT_Float32:
        no_data = -3.4e+38
    elif data_type == gdal.GDT_Float64:
        no_data = -1.7e+308
    else:
        no_data = 0

    assert(raster_array.shape == value_array.shape)

    value_array[raster_array==raster_nodata] = no_data

    return no_data

def array2raster(new_raster_fn, raster_fn, array, data_type=None, unitization=False):
    raster = gdal.Open(raster_fn)
    geotransform = raster.GetGeoTransform()
    origin_x = geotransform[0]
    origin_y = geotransform[3]
    pixel_width = geotransform[1]
    pixel_heigh = geotransform[5]
    cols = array.shape[1]
    rows = array.shape[0]

    print(f"geotransform: {geotransform}")

    if data_type is None:
        data_type = gdal.GDT_Float32

    if unitization:
        array = array * pixel_width

    no_data = setNoDataValue(raster_fn, array, data_type)

    driver = gdal.GetDriverByName('GTiff')
    out_raster = driver.Create(new_raster_fn, cols, rows, 1, data_type)
    out_raster.SetGeoTransform((origin_x, pixel_width, 0, origin_y, 0, pixel_heigh))
    out_band = out_raster.GetRasterBand(1)
    out_band.WriteArray(array)
    out_band.SetNoDataValue(no_data)
    out_raster_srs = osr.SpatialReference()
    out_raster_srs.ImportFromWkt(raster.GetProjectionRef())
    out_raster.SetProjection(out_raster_srs.ExportToWkt())
    out_band.FlushCache()

def main(shape_fn, raster_fn, d_type):
    cost_array = raster2array(raster_fn)
    # print(f"shape: {cost_array.shape}")

    offsets = []
    points = vector2array(shape_fn, raster_fn)
    for point in points:
        offset = coord2pixel_offset(raster_fn, point[0], point[1])
        # find_costs((indexY, indexX))
        offsets.append((offset[1], offset[0]))

    # print(f"offset: {offsets}")
    indices = MCP_Geometric(cost_array)
    cost_surface, back_track = indices.find_costs(offsets)

    array2raster('COST.tif', raster_fn, cost_surface, data_type=d_type, unitization=True)
    array2raster('BACK.tif', raster_fn, back_track, data_type=gdal.GDT_Int16)

if __name__ == "__main__":
    # main('source4326.shp', 'dem4326.tif', 'OUT.tif')
    cost_surface_fn = None
    shape_fn = None
    data_type = None

    argv = sys.argv

    i = 1
    while i < len(argv):
        arg = argv[i]

        if arg == '-t' or arg == '-type':
            i = i + 1
            data_type = argv[i]

        elif shape_fn is None:
            shape_fn = argv[i]

        elif cost_surface_fn is None:
            cost_surface_fn = argv[i]

        else:
            print("python script.py -t float32 shape.shp cost.tif")
            sys.exit(1)

        i = i + 1

    type = {
        'byte': gdal.GDT_Byte,
        'int16': gdal.GDT_Int16,
        'int32': gdal.GDT_Int32,
        'float32': gdal.GDT_Float32,
        'float64': gdal.GDT_Float64
    }

    main(shape_fn, cost_surface_fn, d_type=type.get(data_type, gdal.GDT_Float32))

